from pylab import *

CODE_PATH = '/home/xiaomin/wxm/Code/KaggleCCS'
caffe_root = '/opt/caffe/'  # this file should be run from {caffe_root}/examples (otherwise change this line)
data_root = '/home/xiaomin/wxm/Data/KaggleCCS/'
import sys
import tempfile

sys.path.insert(0, caffe_root + 'python')
import caffe
import os

os.chdir(data_root)
os.chdir('DataForCaffe')

from caffe import layers as L, params as P

caffe.set_device(1)
caffe.set_mode_gpu()

NET_PATH_TEST = CODE_PATH + '/' + 'prototxt/res50/ResNet-50-deploy.prototxt'
NET_PATH_TRAIN = CODE_PATH + '/' + 'prototxt/res50/ResNet50_train_val.prototxt'
WEIGHTS_PATH = '/home/xiaomin/wxm/Data/ResNet/ResNet-50-model.caffemodel'
SOLVER_PATH = CODE_PATH + '/' + 'prototxt/res50/solver.prototxt'


def run_solvers(niter, solvers, disp_interval=10):
    """Run solvers for niter iterations,
       returning the loss and accuracy recorded each iteration.
       `solvers` is a list of (name, solver) tuples."""
    blobs = ('loss', 'acc')
    loss, acc = ({name: np.zeros(niter) for name, _ in solvers}
                 for _ in blobs)
    for it in range(niter):
        for name, s in solvers:
            s.step(1)  # run a single SGD step in Caffe
            loss[name][it], acc[name][it] = (s.net.blobs[b].data.copy()
                                             for b in blobs)
        if it % disp_interval == 0 or it + 1 == niter:
            loss_disp = '; '.join('%s: loss=%.3f, acc=%2d%%' %
                                  (n, loss[n][it], np.round(100 * acc[n][it]))
                                  for n, _ in solvers)
            print '%3d) %s' % (it, loss_disp)
            # Save the learned weights from both nets.
    weight_dir = tempfile.mkdtemp()
    weights = {}
    for name, s in solvers:
        filename = 'weights.%s.caffemodel' % name
        weights[name] = os.path.join(weight_dir, filename)
        s.net.save(weights[name])
    return loss, acc, weights

solver = caffe.get_solver(SOLVER_PATH)
solver.net.copy_from(WEIGHTS_PATH)

scratch_solver = caffe.get_solver(SOLVER_PATH)

niter = 3000  # number of iterations to train

print 'Running solvers for %d iterations...' % niter
solvers = [('pretrained', solver),
           ('scratch', scratch_solver)]
loss, acc, weights = run_solvers(niter, solvers)
print 'Done.'

train_loss, scratch_train_loss = loss['pretrained'], loss['scratch']
train_acc, scratch_train_acc = acc['pretrained'], acc['scratch']
style_weights, scratch_style_weights = weights['pretrained'], weights['scratch']

# Delete solvers to save memory.
del solver, scratch_solver, solvers
